Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Implement specialization constant support in numthreads / local_size #5963

Merged
merged 15 commits into from
Jan 14, 2025
Merged

Implement specialization constant support in numthreads / local_size #5963

merged 15 commits into from
Jan 14, 2025

Conversation

juliusikkala
Copy link
Contributor

@juliusikkala juliusikkala commented Dec 30, 2024 •
edited
Loading

Adds support for using SPIR-V specialization constants in numthreads:

[vk::constant_id(0)] const int WG_SIZE = 8;

[shader("compute")]
[numthreads(1, WG_SIZE, 1)]
void main(uint3 tid : SV_DispatchThreadID)
{
}

Also adds support in GLSL mode:

#version 460

layout(constant_id = 0) const uint size = 8;
// The standard local_size_x_id is supported, but also local_size_x = SPECCONSTNAME.
layout(local_size_x_id = 0, local_size_y = size) in;

void main()
{
}

Similar to the local_size_x_id feature from GLSL, this allows users to generate shader permutations with different work group shapes / thread counts after the shader has already been compiled to SPIR-V.

Currently, this works with the SPIR-V (using LocalSizeId execution mode) and GLSL (using local_size_{x,y,z}_id) targets. I'm not sure if there even is anything that could be done for the other targets, other than to just present diagnostics about this feature not being supported there.

@juliusikkala juliusikkala force-pushed the spec-const-numthreads branch from cde9d98 to 1862056 Compare January 5, 2025 11:59
@juliusikkala juliusikkala force-pushed the spec-const-numthreads branch from 1862056 to 1c7b6cc Compare January 6, 2025 12:55
@juliusikkala
Copy link
Contributor Author

I've now tested this in a real Vulkan program in a couple of different contexts, and this feature seems to now work correctly.

@juliusikkala juliusikkala marked this pull request as ready for review January 6, 2025 13:04
@juliusikkala juliusikkala requested a review from a team as a code owner January 6, 2025 13:04
csyonghe
csyonghe reviewed Jan 8, 2025
View reviewed changes
Copy link
Collaborator

@csyonghe csyonghe left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add a test for this feature?

csyonghe
csyonghe previously approved these changes Jan 9, 2025
View reviewed changes
Copy link
Collaborator

@csyonghe csyonghe left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me. I wish there is a cleaner way to add this feature, but what is implemented in this PR is probably the easiest path forward. If I understand the code correctly, this implementation won't allow something like [numthreads(specConst+1, 1, 2)], right?

source/slang/slang-ast-modifier.h Show resolved Hide resolved
@csyonghe csyonghe added the pr: non-breaking PRs without breaking changes label Jan 9, 2025
@juliusikkala
Copy link
Contributor Author

It currently doesn't allow doing math on the specialization constant, because I'm fairly sure that that can't be expressed in SPIR-V. I'll fix that compiler warning later today, I've been developing on Linux only so I guess MSVC errors are to be expected.

@juliusikkala
Copy link
Contributor Author

juliusikkala commented Jan 12, 2025 •
edited
Loading

Oh no... Sorry for the spam, I made a mistake. I didn't mean to add tangent-vector & expipiplus1 as reviewers, but something went really wrong with merging master into this branch...

Anyway, I fixed one more issue with this PR (thread group sizes in reflection were busted), and re-ran tests locally on a Linux and Windows machine. I think it should be OK now. (previously, all tests were showing "failed" due to #6050)

@juliusikkala juliusikkala force-pushed the spec-const-numthreads branch from 7a2393a to 5544737 Compare January 12, 2025 17:42
@csyonghe csyonghe merged commit cbdc7e1 into shader-slang:master Jan 14, 2025
1 check passed
@fairywreath fairywreath mentioned this pull request Jan 15, 2025
Merged
@juliusikkala juliusikkala mentioned this pull request Jan 17, 2025
Closed
@fknfilewalker
Copy link
Contributor

fknfilewalker commented Jan 19, 2025 •
edited
Loading

It seems that this implementation does not cover setting the wave size via constants e.g.,

[vk::constant_id(0)] const int CONST_VAR = 1;
[WaveSize(CONST_VAR, 1, 1)]

Edit: Ok I am confused, there is also a NumThreads version for which it does not work? What is the difference to the lower case version?

@juliusikkala
Copy link
Contributor Author

juliusikkala commented Jan 19, 2025 •
edited
Loading

@fknfilewalker Is there a platform where you can set wave size / subgroup size with a specialization constant? With a quick scan through the SPIR-V spec, I couldn't find an instruction that would allow for that.

AFAIK numthreads() and NumThreads() should be equivalent (at least they map to the same AST node in core.meta.slang). If there's a difference, that's probably a bug. I don't have the time to look at this for a few days, but it'd help if you can specify in which way it doesn't work. Does it compile, does the SPIR-V pass validation, or does it silently just not do what it should during runtime? What does the SPIR-V disassembly look like for a simple program that fails?

@fknfilewalker
Copy link
Contributor

fknfilewalker commented Jan 19, 2025 •
edited
Loading

With it does not work I meant that setting the size via specialization constants does not compile for NumTheads as well as WaveSize. Explicitly specifying NumTheads or WaveSize does compiles, though WaveSize seems to be doing nothing. I am not sure what WaveSize is for...

@juliusikkala
Copy link
Contributor Author

Huh, odd, I can't replicate that. NumThreads is working just fine for me, identical to numthreads. Can you post a minimal code example of a full slang shader that doesn't compile with NumThreads? The one in the first message of this thread does work for me, with slangc -target spirv test.slang -o test.spv, and produces identical output for both the lowercase and uppercase NumThreads versions.

@fknfilewalker
Copy link
Contributor

Sorry I was wrong, NumThreads works. I mixed things up. It is just WaveSize that does not work, and as we discussed, not sure if it is even useful. Sorry for the troubles.

@juliusikkala
Copy link
Contributor Author

No problem!

Btw, about the WaveSize thing: "Wave" is Subgroup in Khronos terminology, you can read more about that feature here. Some hardware allows you to select how wide the "Wave"/subgroup is (IIRC at least AMD RDNA3 has 32/64 options and Intel Arc GPUs has even more alternatives).

@fknfilewalker
Copy link
Contributor

fknfilewalker commented Jan 19, 2025 •
edited
Loading

Interesting that AMD and Intel can configure this size. I think for SPIRV this is set via the execution mode 35.
https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#Execution_Mode

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@csyonghe csyonghe csyonghe approved these changes

@tangent-vector tangent-vector Awaiting requested review from tangent-vector

@expipiplus1 expipiplus1 Awaiting requested review from expipiplus1

Assignees
No one assigned
Labels
pr: non-breaking PRs without breaking changes
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@juliusikkala @fknfilewalker @csyonghe